#include "imageinterprethelper.h"

#include <curl/curl.h>
#include <jsoncpp/json/json.h>

#include <iostream>

#include "utils/utils.h"

namespace {
size_t WriteCallback(void* contents, size_t size, size_t nmemb, std::string* userp)
{
    size_t totalSize = size * nmemb;
    userp->append((char*)contents, totalSize);
    return totalSize;
}

std::string extractStringBetweenBraces(const std::string& inputStr) {
    size_t startPos = inputStr.find('{');
    if (startPos == std::string::npos) {
        return std::string(); // 未找到左大括号
    }

    size_t endPos = inputStr.find('}', startPos);
    if (endPos == std::string::npos) {
        return std::string(); // 未找到右大括号
    }

    return inputStr.substr(startPos, endPos - startPos + 1);
}

} // namespace

std::string ImageInterpretHelper::getImageSummary(const std::string &imagePath) const
{
    const std::string url = "http://10.41.116.101:57693/minicpm_inference";

    const std::string userQuery = "Generate a summary for this image";
    const std::string sceneName = "image_abstract";
    const int maxLength = 200;

    std::string base64Data = utils::encodeImageToBase64(imagePath);

    std::string postJson = convertRequestJson(base64Data, userQuery, sceneName, maxLength);

    std::string responseJson = postRequest(url, postJson);

    return extractSummaryFromResponse(responseJson);
}

std::vector<std::string> ImageInterpretHelper::getImageTags(const std::string &imagePath) const
{
    const std::string url = "http://10.41.116.101:57693/minicpm_inference";

    const std::string userQuery = "为这张图片生成标签";
    const std::string sceneName = "image_label_extraction";
    const int maxTags = 20;
    const int maxTagLength = 15;

    std::string base64Data= utils::encodeImageToBase64(imagePath);

    std::string postJson = convertRequestJson(base64Data, userQuery, sceneName, 0, maxTags, maxTagLength);

    std::string responseJson = postRequest(url, postJson);

    return extractTagsFromResponse(responseJson);
}

std::string ImageInterpretHelper::postRequest(const std::string &url, const std::string &data) const
{
    CURL* curl;
    CURLcode res;
    std::string readBuffer;

    curl_global_init(CURL_GLOBAL_DEFAULT);
    curl = curl_easy_init();
    if (!curl) {
        std::cerr << "curl_easy_init() failed" << std::endl;
        curl_global_cleanup();
        return std::string();
    }

    curl_easy_setopt(curl, CURLOPT_URL, url.c_str());
    curl_easy_setopt(curl, CURLOPT_POSTFIELDS, data.c_str());
    curl_easy_setopt(curl, CURLOPT_WRITEFUNCTION, WriteCallback);
    curl_easy_setopt(curl, CURLOPT_WRITEDATA, &readBuffer);

    res = curl_easy_perform(curl);
    if(res != CURLE_OK) {
        std::cerr << "curl_easy_perform() failed: " << curl_easy_strerror(res);
    }

    curl_easy_cleanup(curl);
    curl_global_cleanup();

    return readBuffer;
}

std::string ImageInterpretHelper::convertRequestJson(const std::string &base64Image,
                                                     const std::string &userQuery,
                                                     const std::string &sceneName,
                                                     int maxLength,
                                                     int labelsNum,
                                                     int labelsLength) const
{
    Json::Value root;
    root["image_base64"] = base64Image;
    root["user_query"] = userQuery;
    root["scene_name"] = sceneName;
    if (maxLength > 0)
        root["max_length"] = std::to_string(maxLength);
    if (labelsNum > 0)
        root["labels_num"] = std::to_string(labelsNum);
    if (labelsLength > 0)
        root["labels_length"] = std::to_string(labelsLength);

    Json::StreamWriterBuilder builder;
    builder["indentation"] = "";
    return Json::writeString(builder, root);
}

std::string ImageInterpretHelper::extractSummaryFromResponse(const std::string &responseJson) const
{
    Json::Value root;
    Json::Reader reader;
    if (!reader.parse(responseJson, root)) {
        std::cerr << "Failed to parse JSON response: " << reader.getFormattedErrorMessages();
        return std::string();
    }
    if (!root.isMember("modelResponse")) {
        std::cerr << "Response JSON does not contain 'modelResponse' key";
        return std::string();
    }

    Json::Value subNode;
    std::string subJsonString = root["modelResponse"].asString();
    subJsonString = extractStringBetweenBraces(subJsonString);

    if (!reader.parse(subJsonString, subNode)) {
        std::cerr << "Failed to parse sub JSON response: " << reader.getFormattedErrorMessages();
        return std::string();
    }
    if (!subNode.isMember("imageAbstract")) {
        std::cerr << "Sub JSON response does not contain 'imageAbstract' key";
        return std::string();
    }

    return subNode["imageAbstract"].asString();
}


std::vector<std::string> ImageInterpretHelper::extractTagsFromResponse(const std::string &responseJson) const
{
    Json::Value root;
    Json::Reader reader;
    if (!reader.parse(responseJson, root)) {
        std::cerr << "Failed to parse JSON response: " << reader.getFormattedErrorMessages();
        return std::vector<std::string>();
    }
    if (!root.isMember("modelResponse")) {
        std::cerr << "Response JSON does not contain 'modelResponse' key";
        return std::vector<std::string>();
    }

    Json::Value subNode;
    std::string subJsonString = root["modelResponse"].asString();
    if (!reader.parse(subJsonString, subNode)) {
        std::cerr << "Failed to parse sub JSON response: " << reader.getFormattedErrorMessages();
        return std::vector<std::string>();
    }
    if (!subNode.isMember("presetLabels")) {
        std::cerr << "Response JSON does not contain 'presetLabels' key" << std::endl;
        return std::vector<std::string>();
    }
    if (!subNode.isMember("freeLabels")) {
        std::cerr << "Response JSON does not contain 'freeLabels' key" << std::endl;
        return std::vector<std::string>();
    }

    const Json::Value &presetLabels = subNode["presetLabels"];
    const Json::Value &freeLabels = subNode["freeLabels"];
    if (!presetLabels.isArray()) {
        std::cerr << "Response JSON does not contain 'presetLabels' array";
        return std::vector<std::string>();
    }
    if (!freeLabels.isArray()) {
        std::cerr << "Response JSON does not contain 'freeLabels' array";
        return std::vector<std::string>();
    }

    std::vector<std::string> tags;
    for (const auto &label : presetLabels) {
        if (label.isMember("label"))
            tags.push_back(label["label"].asString());
    }
    for (const auto &label : freeLabels) {
        if (label.isMember("label"))
            tags.push_back(label["label"].asString());
    }

    return tags;
}
